import json

from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.utils.json import parse_json_markdown


def stateParse(gState):
    print("startParse--------------")
    tool_response_prompt = """工具响应：
------
{tools_response}
------

请根据工具的响应判断，是否能够回答问题：

{input}

请根据工具响应的内容，思考接下来回复。回复格式严格按照前面所说的2种JSON回复格式，选择其中1种进行回复。请记住只选择单个选项格式，以JSON格式回复内容，不要回复其他内容

"""
    if len(gState["messages"]) > 1:
        chat_history = [
            HumanMessage(content=gState["messages"][0]),
            AIMessage(json.dumps(gState["messages"][-1])),
        ]
        agent_scratchpad = [
            AIMessage(json.dumps(gState["messages"][-2])),
            HumanMessage(
                content=tool_response_prompt.format(
                    input=gState["messages"][0],
                    tools_response=json.dumps(gState["messages"][-1])
                ),
            )
        ]
    else:
        agent_scratchpad = []
    return {
        "input": gState["messages"][0],
        "chat_history": gState["chat_history"],
        "agent_scratchpad": agent_scratchpad
    }


def msgParse(message):
    message = message.content.replace("'", '"')
    message = parse_json_markdown(message)

    return {"messages": [message]}